using System;
using System.Collections.Generic;
using System.Text;

namespace LSO
{
    static class LSOOptimizer
    {
        /// <summary>
        /// Mini-assembler just for sequences of instructions
        /// </summary>
        /// <param name="instructions"></param>
        /// <returns></returns>
        public static byte[] AssembleInstructions(List<LSO.Instruction> instructions)
        {
            int length = 0;
            foreach (LSO.Instruction instruction in instructions)
                length += instruction.Size;
            byte[] b = new byte[length];
            int offset = 0;
            byte[] buffer;
            foreach (LSO.Instruction instruction in instructions)
            {
                buffer = instruction.ToBytes();
                Buffer.BlockCopy(buffer, 0, b, offset, buffer.Length);
                offset += buffer.Length;
            }
            return b;
        }

        /// <summary>
        /// Compares two instructions including their arguments.
        /// Not sure why this is necessary
        /// </summary>
        public static bool SameInstruction(LSO.Instruction first, LSO.Instruction second)
        {
            if (first.Info.Opcode != second.Info.Opcode)
                return false;
            if (first.Info.NumberOfArgs != second.Info.NumberOfArgs)
                return false;
            int argcount = first.Info.NumberOfArgs;
            for (int i = 0; i < argcount; i++)
            {
                if ((first.Arguments[i].Equals(second.Arguments[i])) == false)
                    return false;
            }
            return true;
        }

        /// <summary>
        /// Returns true if no jumps land within the defined region of instructions
        /// (Greater than offset, less than offset + length)
        /// </summary>
        public static bool OkayToReplace(List<LSO.Instruction> haystack, int offset, int length)
        {
            List<int> Offsets = new List<int>();
            int byteoffset = 0;
            for (int i = 0; i < haystack.Count; i++)
            {
                Offsets.Add(byteoffset);
                byteoffset += haystack[i].Size;
            }

            int bytestart = Offsets[offset];
            int byteend = Offsets[offset + length];

            LSO.Instruction instruction;
            int jump;
            for (int i = 0; i < haystack.Count; i++)
            {
                instruction = haystack[i];
                jump = 0;
                if (instruction.Info.Official == "JUMP")
                    jump = Convert.ToInt32(instruction.Arguments[0]) + instruction.Size;
                else if (instruction.Info.Official == "JUMPIF" || instruction.Info.Official == "JUMPNIF")
                    jump = Convert.ToInt32(instruction.Arguments[1]) + instruction.Size;
                if (jump != 0)
                {
                    jump = Offsets[i] + jump;
                    if (jump > bytestart && jump < byteend)
                        return false;
                }
            }
            return true;
        }

        public static int SizeDifference(List<LSO.Instruction> from, List<LSO.Instruction> to)
        {
            int fromsize = 0;
            foreach (LSO.Instruction inst in from)
                fromsize += inst.Size;
            int tosize = 0;
            foreach (LSO.Instruction inst in to)
                tosize += inst.Size;
            return fromsize - tosize;
        }

        private static List<LSO.Instruction> RecalculateJumps(List<LSO.Instruction> instructions, int median, int bytedifference)
        {
            List<int> Offsets = new List<int>();
            int byteoffset = 0;
            for (int i = 0; i < instructions.Count; i++)
            {
                Offsets.Add(byteoffset);
                byteoffset += instructions[i].Size;
            }

            int bytemedian = Offsets[median];

            LSO.Instruction instruction;
            int jump;
            for (int i = 0; i < instructions.Count; i++)
            {
                jump = 0;
                instruction = instructions[i];
                if (instruction.Info.Official == "JUMP")
                    jump = Convert.ToInt32(instruction.Arguments[0]) + instruction.Size;
                else if (instruction.Info.Official == "JUMPIF" || instruction.Info.Official == "JUMPNIF")
                    jump = Convert.ToInt32(instruction.Arguments[1]) + instruction.Size;
                if (jump != 0)
                {
                    jump = Offsets[i] + jump;
                    if (Offsets[i] < bytemedian)
                    {
                        if (jump >= bytemedian)
                            jump -= bytedifference;
                    }
                    else
                    {
                        if (jump < bytemedian)
                            jump += bytedifference;
                    }
                    jump = jump - Offsets[i] - instruction.Size;
                    if (instruction.Info.Official == "JUMP")
                        instruction.Arguments[0] = jump;
                    else
                        instruction.Arguments[1] = jump;
                }
            }
            return instructions;
        }

        public static class FindAndReplaceOptimizer
        {
            private static Dictionary<List<LSO.Instruction>, List<LSO.Instruction>> LookupTable = new Dictionary<List<LSO.Instruction>, List<LSO.Instruction>>();

            private static int ListFind(List<LSO.Instruction> haystack, List<LSO.Instruction> needle, int offset)
            {
                int si;
                int next = offset;
                for (int di = offset; di < haystack.Count; di++)
                {
                    si = 0;
                    di = next;
                    next = di + 1;
                    while (SameInstruction(haystack[di], needle[si]))
                    {
                        si++;
                        di++;
                        if (si == needle.Count)
                            return next - 1;
                        if (di == haystack.Count)
                            break;
                    }
                }
                return -1;
            }

            private static List<LSO.Instruction> ListReplace(List<LSO.Instruction> haystack, int replacedcount, List<LSO.Instruction> needle, int offset)
            {
                List<LSO.Instruction> output = new List<LSO.Instruction>();
                for (int i = 0; i < offset; i++)
                    output.Add(haystack[i]);
                foreach(LSO.Instruction inst in needle)
                    output.Add(inst);
                for(int i = offset + replacedcount; i < haystack.Count; i++)
                {
                    output.Add(haystack[i]);
                }
                return output;
            }

            private static void LUTAdd(string find, string replace)
            {
                List<LSO.Instruction> findinst = LSODisassembler.Disassemble(Helpers.Hex2Bytes(find));
                List<LSO.Instruction> replaceinst = LSODisassembler.Disassemble(Helpers.Hex2Bytes(replace));
                LookupTable.Add(findinst, replaceinst);
            }

            static FindAndReplaceOptimizer()
            {
                #region new
                // pushargi 1 add integer, integer -> bitnot neg integer
                // thx Strife
                LUTAdd("5e000000017011", "818001");
                // pushargi 1 sub integer, integer -> neg integer bitnot
                // thx Strife
                LUTAdd("5e000000017111", "800181");
                #endregion

                #region POPARG *
                // poparg 0 -> nothing
                LUTAdd("0600000000", "");
                // poparg 4 -> pop
                LUTAdd("0600000004", "01");
                // poparg 8 -> pop pop
                LUTAdd("0600000008", "0101");
                // poparg 12 -> popv
                LUTAdd("060000000c", "04");
                // poparg 16 -> popq
                LUTAdd("0600000010", "05");
                // poparg 20 -> popq pop
                LUTAdd("0600000014", "0501");
                // poparg 24 -> popv popv
                LUTAdd("0600000018", "0404");
                // poparg 28 -> popq popv
                LUTAdd("060000001c", "0504");
                // poparg 32 -> popq popq
                LUTAdd("0600000020", "0505");
                // poparg 36 -> popv popv popv
                LUTAdd("0600000024", "040404");
                // poparg 40 -> popq popv popv
                LUTAdd("0600000028", "050404");
                // poparg 44 -> popq popq popv
                LUTAdd("060000002c", "050504");
                // poparg 48 -> popq popq popq
                LUTAdd("0600000030", "050505");
                // poparg 52 -> popq popq popq pop
                LUTAdd("0600000034", "05050501");
                // poparg 56 -> popq popq popv popv
                LUTAdd("0600000038", "05050404");
                // poparg 60 -> popq popq popq popv
                LUTAdd("060000003c", "05050504");
                // poparg 64 -> popq popq popq popq
                LUTAdd("0600000040", "05050505");
                #endregion

                #region POP POP POP ...
                // pop pop pop pop -> popq
                LUTAdd("01010101", "05");
                // pop pop pop -> popv
                LUTAdd("010101", "04");
                #endregion

                #region PUSHARGE *
                // pusharge 0 -> nothing
                LUTAdd("6600000000", "");
                // pusharge 4 -> pushe
                LUTAdd("6600000004", "63");
                // pusharge 8 -> pushe pushe
                LUTAdd("6600000008", "6363");
                // pusharge 12 -> pushev
                LUTAdd("660000000c", "64");
                // pusharge 16 -> pusheq
                LUTAdd("6600000010", "65");
                // pusharge 20 -> pusheq pushe
                LUTAdd("6600000014", "6563");
                // pusharge 24 -> pushev pushev
                LUTAdd("6600000018", "6464");
                // pusharge 28 -> pusheq pushev
                LUTAdd("660000001c", "6564");
                // pusharge 32 -> pusheq pusheq
                LUTAdd("6600000020", "6565");
                // pusharge 36 -> pushev pushev pushev
                LUTAdd("6600000024", "646464");
                // pusharge 40 -> pusheq pushev pushev
                LUTAdd("6600000028", "656464");
                // pusharge 44 -> pusheq pusheq pushev
                LUTAdd("660000002c", "656564");
                // pusharge 48 -> pusheq pusheq pusheq
                LUTAdd("6600000030", "656565");
                // pusharge 52 -> pusheq pushev pushev pushev
                LUTAdd("6600000034", "65646464");
                // pusharge 56 -> pusheq pusheq pushev pushev
                LUTAdd("6600000038", "65656464");
                // pusharge 60 -> pusheq pusheq pusheq pushev
                LUTAdd("660000003c", "65656564");
                // pusharge 64 -> pusheq pusheq pusheq pusheq
                LUTAdd("6600000040", "65656565");
                #endregion

                #region PUSHARG* 0
                // pushargi 0 -> pushe
                LUTAdd("5e00000000", "63");
                // pushargf 0
                LUTAdd("5f00000000", "63");
                // pushargv <0, 0, 0> -> pushev
                LUTAdd("61000000000000000000000000", "64");
                // pushargq <0, 0, 0, 0> -> pusheq
                LUTAdd("6200000000000000000000000000000000", "65");
                #endregion

                #region ADD/SUB 0
                // pushe add integer, integer -> nothing
                LUTAdd("637011", "");
                // pushe sub integer, integer -> nothing
                LUTAdd("637111", "");
                #endregion

                #region These are smaller, but there's a chance they may be slower
                // pushargi 1 -> pushe boolnot
                LUTAdd("5E00000001", "6382");
                // pushargi -1 -> pushe bitnot
                LUTAdd("5EFFFFFFFF", "6381");
                #endregion

                #region PUSHE PUSHE PUSHE ...
                // pushe pushe pushe pushe -> pusheq
                LUTAdd("63636363", "65");
                // pushe pushe pushe -> pushev
                LUTAdd("636363", "64");
                #endregion

                #region ideas
                // push 0 push 4 push 8 -> pushv 0 ?
                // pushargi 2 -> pushe boolnot dup shl ?
                // callib_two_byte 0x00** -> calllib 0x** ?
                // push(g) x push(g) x pushargi 1 add/sub int, int store(g) x pop pop
                // -> push(g) x pushargi 1 add/sub int, int store(g) x pop (Strife)
                #endregion
            }

            public static List<LSO.Instruction> Optimize(List<LSO.Instruction> instructions)
            {
                int offset;
                foreach (KeyValuePair<List<LSO.Instruction>, List<LSO.Instruction>> kvp in LookupTable)
                {
                    offset = 0;
                    do
                    {
                        offset = ListFind(instructions, kvp.Key, offset);
                        if (offset != -1)
                        {
                            if (OkayToReplace(instructions, offset, kvp.Key.Count))
                            {
                                instructions = ListReplace(instructions, kvp.Key.Count, kvp.Value, offset);
                                instructions = RecalculateJumps(instructions, offset + kvp.Value.Count, SizeDifference(kvp.Key, kvp.Value));
                            }
                            else
                                offset++;
                        }
                    } while (offset != -1);
                }
                return instructions;
            }

            public static LSOAssembly PerformOptimizations(LSOAssembly assembly)
            {
                List<LSO.Instruction> instructions;
                foreach (LSOAssembly.FunctionCodeChunk chunk in assembly.FunctionCodeChunks)
                {
                    instructions = LSODisassembler.Disassemble(chunk.Code);
                    instructions = Optimize(instructions);
                    chunk.Code = AssembleInstructions(instructions);
                }
                foreach (LSOAssembly.EventHandlerCodeChunk chunk in assembly.EventHandlerCodeChunks)
                {
                    instructions = LSODisassembler.Disassemble(chunk.Code);
                    instructions = Optimize(instructions);
                    chunk.Code = AssembleInstructions(instructions);
                }
                return assembly;
            }
        }

        public static class HeaderStringOptimizer
        {
            public static LSOAssembly PerformOptimizations(LSOAssembly assembly)
            {
                foreach (LSOAssembly.StaticBlock block in assembly.StaticBlocks)
                    block.Header.StringData = "";
                foreach (LSOAssembly.StateBlock block in assembly.StateBlocks)
                    block.Header.StringData = "";
                foreach (LSOAssembly.EventHandlerCodeChunk block in assembly.EventHandlerCodeChunks)
                    block.Header.StringData = "";
                return assembly;
            }
        }

        public static LSOAssembly PerformAllOptimizations(LSOAssembly assembly)
        {
            assembly = FindAndReplaceOptimizer.PerformOptimizations(assembly);
            assembly = HeaderStringOptimizer.PerformOptimizations(assembly);
            return assembly;
        }
    }
}
